package be.fedict.eid.diagnostics.system.windows.jna;

import be.fedict.eid.diagnostics.system.windows.proxies.Advapi32;
import be.fedict.eid.diagnostics.system.windows.proxies.Kernel32;
import be.fedict.eid.diagnostics.system.windows.proxies.Advapi32.LUID;
import be.fedict.eid.diagnostics.system.windows.proxies.Advapi32.LUID_AND_ATTRIBUTES;
import be.fedict.eid.diagnostics.system.windows.proxies.Advapi32.TOKEN_PRIVILEGES;

import com.sun.jna.Native;
import com.sun.jna.platform.win32.WinDef.DWORD;
import com.sun.jna.platform.win32.WinNT.HANDLE;
import com.sun.jna.platform.win32.WinNT.HANDLEByReference;

public class SecurityHandler {
	
	private static final int TOKEN_ADJUST_PRIVILEGES = 0x00000020;
	private static final int TOKEN_QUERY = 0x00000008;
	private static final String SE_DEBUG_NAME = "SeDebugPrivilege";
	private static final int SE_PRIVILEGE_ENABLED = 0x00000002;
	private static final int ERROR_NOT_ALL_ASSIGNED = 0x514; // 1300

	private static boolean hasDebug = false;
	
	private static boolean SetPrivilege(HANDLEByReference hToken, 
								String lpszPrivilege, 
								boolean bEnablePrivilege ) {
		
		Advapi32 advapi32 = Advapi32.INSTANCE;
		
		TOKEN_PRIVILEGES tp = new Advapi32.TOKEN_PRIVILEGES(new DWORD(1));
		LUID luid = new LUID();
		if (! advapi32.LookupPrivilegeValueA(null, lpszPrivilege, luid)) {
			return false;
		}
		tp.Privileges[0] = new LUID_AND_ATTRIBUTES();
		tp.Privileges[0].Luid = luid;
		if (bEnablePrivilege)
			tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
		else
			tp.Privileges[0].Attributes = 0;
		
		if(! advapi32.AdjustTokenPrivileges(
				hToken.getValue(),
				false,
				tp,
				tp.size(),
				null,
				null)) {
			return false;
			
		}
		
		if (Native.getLastError() == ERROR_NOT_ALL_ASSIGNED) {
			return false;
		}
		
		hasDebug = true;
		return true;
		
	}
	
	public static boolean getDebugPrivileges() {
		Advapi32 advapi32 = Advapi32.INSTANCE;
		
		HANDLE hCurrentProcess = Kernel32.INSTANCE.GetCurrentProcess();
		HANDLEByReference hToken = new HANDLEByReference();
		
		// retrieve process privilege tokens
		boolean success = advapi32.OpenProcessToken(
										hCurrentProcess,
										TOKEN_ADJUST_PRIVILEGES | TOKEN_QUERY,
										hToken);
		
		if (success) {
			return SetPrivilege(hToken, SE_DEBUG_NAME, true);
		}
		else {
			return false;
		}
	}

	public static boolean hasDebugPrivileges() {
		return hasDebug;
	}
}
